{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "798471a7-188f-42c5-ab56-55ffdb721c9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Goes with the runCommandsOnNewCallback.yml file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d966d1f-7d2c-46db-877e-c356f31c57eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from mythic import mythic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b34fd529-79cd-4a5a-875a-c2c0c9c4ef18",
   "metadata": {},
   "outputs": [],
   "source": [
    "mythic_instance = await mythic.login(\n",
    "        username=\"mythic_admin\",\n",
    "        password=\"mythic_password\",\n",
    "        server_ip=\"mythic_nginx\",\n",
    "        server_port=7443,\n",
    "        timeout=-1\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f697ff41-9b19-4a1d-8444-581441ccd56d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ################ attributes available for env. in trigger:callback_new ################\n",
    "## these are the fields you can access via env. in your Inputs. For example:\n",
    "## env.display_id <-- that would give you the triggering callback's display_id field\n",
    "callback_attribute_query = \"\"\"\n",
    "query MyQuery {\n",
    "  callback(limit: 1) {\n",
    "    active\n",
    "    agent_callback_id\n",
    "    architecture\n",
    "    crypto_type\n",
    "    current_time\n",
    "    dead\n",
    "    dec_key\n",
    "    dec_key_base64\n",
    "    description\n",
    "    display_id\n",
    "    domain\n",
    "    enc_key\n",
    "    enc_key_base64\n",
    "    eventstepinstance_id\n",
    "    external_ip\n",
    "    extra_info\n",
    "    host\n",
    "    id\n",
    "    init_callback\n",
    "    integrity_level\n",
    "    ip\n",
    "    last_checkin\n",
    "    locked\n",
    "    locked_operator_id\n",
    "    mythictree_groups\n",
    "    mythictree_groups_string\n",
    "    operation_id\n",
    "    operator_id\n",
    "    os\n",
    "    pid\n",
    "    process_name\n",
    "    process_short_name\n",
    "    registered_payload_id\n",
    "    sleep_info\n",
    "    timestamp\n",
    "    user\n",
    "  }\n",
    "}\n",
    "\"\"\"\n",
    "fields = await mythic.execute_custom_query(\n",
    "    mythic=mythic_instance,\n",
    "    query=callback_attribute_query\n",
    ")\n",
    "print(fields)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "068a8bdd-9154-4a8c-850e-d5aa55979b51",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Trigger Data\n",
    "\n",
    "#trigger_data:\n",
    "#  payload_types:\n",
    "#    - poseidon\n",
    "# This is how you limit this trigger to only happen on callbacks based on certain payload types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9403bab-8655-4292-87cd-1505b5466ede",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Keywords\n",
    "\n",
    "#keywords:\n",
    "#  - poseidon_callback\n",
    "\n",
    "# keywords allow you to initiate a workflow by keyword and custom data instead of waiting for an event to happen\n",
    "# in this case, if we use the keyword `poseidon_callback`, then we can trigger this workflow\n",
    "# there's no requirement for these to be unique though, so you could trigger a bunch with a single keyword\n",
    "# this triggering is available via the Scripting component:\n",
    "trigger_keyword = \"\"\"\n",
    "mutation TriggerByKeyword($keyword: String!, $keywordEnvData: jsonb!){\n",
    "    eventingTriggerKeyword(keyword: $keyword, keywordEnvData: $keywordEnvData){\n",
    "        status\n",
    "        error\n",
    "    }\n",
    "}\n",
    "\"\"\"\n",
    "triggerStatus = await mythic.execute_custom_query(\n",
    "    mythic=mythic_instance,\n",
    "    query=trigger_keyword,\n",
    "    variables={\"keyword\": \"poseidon_callback\", \"keywordEnvData\": {\"display_id\": 1}}\n",
    ")\n",
    "# In this case, we are seeding this workflow's `display_id` to be 1, and providing no other data\n",
    "# You can provide whatever data you want, but you should make sure that you provide everything that's needed by all your steps\n",
    "print(triggerStatus)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72a5522e-b1ad-45c6-b41b-6d35960fd61f",
   "metadata": {},
   "outputs": [],
   "source": [
    "### action: task_create (string inputs)\n",
    "\n",
    "# This works largely the same as how we issue tasks in scripting too, just in yaml form\n",
    "#- name: \"issue whoami\"\n",
    "#    description:\n",
    "#    inputs:\n",
    "#      CALLBACK_ID: env.display_id <-- notice this gets the display_id of the callback that triggered this event\n",
    "#      COMMAND: shell <-- this could be hardcoded in the action_data, but left here as a simple example\n",
    "#    action: task_create\n",
    "#    action_data: <-- the actual data associated with this action, task_create\n",
    "#      callback_display_id: CALLBACK_ID <-- this gets replaced based on our inputs\n",
    "#      params: whoami\n",
    "#      command_name: COMMAND <-- this gets replaced based on our inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69203402-82e2-4ea3-87d0-12dc3dfd825b",
   "metadata": {},
   "outputs": [],
   "source": [
    "### action: task_create (but with a continue_on_error)\n",
    "\n",
    "#  - name: \"getenv\"\n",
    "#    description:\n",
    "#    inputs:\n",
    "#      CALLBACK_ID: env.display_id\n",
    "#    action: task_create\n",
    "#    action_data:\n",
    "#      callback_display_id: CALLBACK_ID\n",
    "#      params: \n",
    "#      command_name: getenv\n",
    "#    continue_on_error: true <-- normally, if one step errors, then the entire workflow is cancelled\n",
    "#                       ^ this is how you can say that if this step errors, still continue on with the rest of the workflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2fdf637-4218-4c06-9760-870fdb9129ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "### action: task_create (with workflow files and dictionaries)\n",
    "\n",
    "#  - name: \"import script\"\n",
    "#    description: \"import the HealthInspector script into the callback\"\n",
    "#    inputs:\n",
    "#      CALLBACK_ID: env.display_id <-- env.display_id gets the display_id from the triggering callback\n",
    "#      FILE_ID: workflow.HealthInspector.js <-- workflow.* gets the agent_file_id associated with files uploaded as part of the workflow\n",
    "#               ^ click the paperclip icon when viewing a workflow to upload / manage files to use with it\n",
    "#    action: task_create\n",
    "#    action_data:\n",
    "#      callback_display_id: CALLBACK_ID <-- repalced based on our inputs\n",
    "#      params_dictionary: <-- notice this is params_dictionary instead of just params, now we can provide non-ambiguous parameters \n",
    "#        file_id: FILE_ID <-- replaced based on our inputs\n",
    "#      command_name: jsimport"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddc5d3dd-ac66-41b2-872b-461db0e135e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "### action: task_create (depends_on and outputs)\n",
    "\n",
    "# - name: \"run script\"\n",
    "#    description: \"run the HealthInspector script\"\n",
    "#    inputs:\n",
    "#      CALLBACK_ID: env.display_id\n",
    "#    action: task_create\n",
    "#    action_data:\n",
    "#      callback_display_id: CALLBACK_ID\n",
    "#      params_dictionary:\n",
    "#        filename: HealthInspector.js <-- this specific command looks up files by name instead of agent_file_id, so we can just specify the name\n",
    "#        code: All_Checks(); <-- the code we want to run\n",
    "#      command_name: jsimport_call\n",
    "#    depends_on:\n",
    "#      - import script <-- the \"name\" field from the task that imported the file. This task won't run until that task executes successfully\n",
    "#    outputs:\n",
    "#      SCRIPT_TASK_ID: id <-- this step created a task, so we can save off any attribute of the task as output, all options shown below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b16c4d62-810c-452a-b0f5-0e3f33bc1502",
   "metadata": {},
   "outputs": [],
   "source": [
    "task_attribute_query = \"\"\"\n",
    "query MyQuery {\n",
    "  task(limit: 1) {\n",
    "    agent_task_id\n",
    "    apitokens_id\n",
    "    callback_id\n",
    "    command_id\n",
    "    command_name\n",
    "    comment\n",
    "    comment_operator_id\n",
    "    completed\n",
    "    completed_callback_function\n",
    "    completed_callback_function_completed\n",
    "    display_id\n",
    "    display_params\n",
    "    eventstepinstance_id\n",
    "    group_callback_function\n",
    "    group_callback_function_completed\n",
    "    has_intercepted_response\n",
    "    id\n",
    "    interactive_task_type\n",
    "    is_interactive_task\n",
    "    operation_id\n",
    "    operator_id\n",
    "    opsec_post_blocked\n",
    "    opsec_post_bypass_role\n",
    "    opsec_post_bypass_user_id\n",
    "    opsec_post_bypassed\n",
    "    opsec_post_message\n",
    "    opsec_pre_blocked\n",
    "    opsec_pre_bypass_role\n",
    "    opsec_pre_bypass_user_id\n",
    "    opsec_pre_bypassed\n",
    "    opsec_pre_message\n",
    "    original_params\n",
    "    parameter_group_name\n",
    "    params\n",
    "    parent_task_id\n",
    "    response_count\n",
    "    status\n",
    "    status_timestamp_preprocessing\n",
    "    status_timestamp_processed\n",
    "    status_timestamp_processing\n",
    "    status_timestamp_submitted\n",
    "    stderr\n",
    "    stdout\n",
    "    subtask_callback_function\n",
    "    subtask_callback_function_completed\n",
    "    subtask_group_name\n",
    "    tasking_location\n",
    "    timestamp\n",
    "    token_id\n",
    "  }\n",
    "}\n",
    "\"\"\"\n",
    "fields = await mythic.execute_custom_query(\n",
    "    mythic=mythic_instance,\n",
    "    query=task_attribute_query\n",
    ")\n",
    "print(fields)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd35df63-efc1-4d57-9463-b52a6d088543",
   "metadata": {},
   "outputs": [],
   "source": [
    "### action: custom_function (apitokens, custom functions, and inputs from other tasks)\n",
    "\n",
    "#  - name: \"check for EDR\"\n",
    "#    description: \"Process the output of the HealthInspector task to see if there's any EDR running\"\n",
    "#    inputs:\n",
    "#      SCRIPT_TASK_ID: \"run script.SCRIPT_TASK_ID\" <-- one of our inputs will be this named output from a previous step\n",
    "#      APIToken: mythic.apitoken <-- we can ask Mythic for a custom APIToken that'll track what we do to this specific step\n",
    "#      CALLBACK_ID: env.display_id <-- still get the same triggering callback's display_id\n",
    "#    action: custom_function\n",
    "#    action_data:\n",
    "#      function_name: HealthInspectorOutputEDRCheck <-- the name of the CustomFunction entry within our container to call\n",
    "#      container_name: opsecChecker <-- the name of the custom container we created that has the above function for us to call\n",
    "#    depends_on:\n",
    "#      - run script <-- the \"name\" field from the other step that needs to complete first"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "972397ae-4a07-41d7-abc5-f1fed9a0c18b",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Other thoughts\n",
    "# If you just load the runCommandsOnNewCallback.yml into Mythic you'll see that there's a red warning about missing the specified container\n",
    "# Mythic won't currently alert you that you specify a workflow file that isn't upload (future update)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
